from __future__ import annotations
from dataclasses import dataclass, field
from codesecurity.feature.objects import AstNode,AstEdge,CommonFeatureSet,Ast
from codesecurity.feature.property_graph import get_anchor_node_js,list_fusion_node_state
from codesecurity.feature.ast_image import AstImageBuilder

from collections import Counter

@dataclass
class FusionNode:
    nodes:list[AstNode]=field(default_factory=list)
    node_value:str=""
    node_type:str=""
    
    @property
    def is_leave(self):
        return True
    
@dataclass 
class FusionAst:
    fusion_nodes:list[FusionNode]=field(default_factory=list)
    edges:list[AstEdge]=field(default_factory=list)

class FusionAstBuilder:
    def __init__(self,anchor_visions:dict) -> None:
        self.anchor_visions=anchor_visions
        self.anchor_type=list(anchor_visions.keys())
        
    def get_merge_state(self,features:CommonFeatureSet,image_builder:AstImageBuilder):
        fusion_dict=self.list_fusion_nodes(features,image_builder)
        
        fusion_nodes=[]
        fusion_mapping={}
        
        for fusion_name in fusion_dict:
            fusion_indexs=fusion_dict[fusion_name]
            fusion_node=self.fusion_node(features,fusion_indexs)
            for node_index in fusion_indexs:
                fusion_mapping[node_index]=fusion_node
            fusion_nodes.append(fusion_node)
    
        return fusion_nodes,fusion_mapping

    def build(self,features:CommonFeatureSet):

        image_builder=AstImageBuilder(features.ast_object)

        fusion_nodes,fusion_mapping=self.get_merge_state(features,image_builder)

        ast=self.merge_nodes(features.ast_object,fusion_mapping)
        
        return ast


    def merge_nodes(self,ast_object:Ast,fusion_mapping:dict):
        new_nodes=[]
        new_edges=[]

        node_mapping_flag={e:False for e in fusion_mapping}

        for i,node in enumerate(ast_object.nodes):
            if i in fusion_mapping:
                if node_mapping_flag[i]==False:
                    new_nodes.append(fusion_mapping[i])
                    node_mapping_flag[i]=True
            else:
                new_nodes.append(node)

        new_index_mapping={e:i for i,e in enumerate(new_nodes)}

        for i,edge in enumerate(ast_object.edges):
            start,end=edge.start_index,edge.end_index

            start=fusion_mapping.get(start,ast_object.nodes[start])
            end=fusion_mapping.get(end,ast_object.nodes[end])

            start=new_index_mapping[start]
            end=new_index_mapping[end]

            new_edges.append(AstEdge(start,end))
        
        return Ast(new_nodes,new_edges)


    def fusion_node(self,features:CommonFeatureSet,fusion_indexes:list[int]):
        core_node=features.nodes[fusion_indexes[0]]
        
        node_type=core_node.ast_type
        node_value=core_node.ast_value
        nodes=[features.nodes[i] for i in fusion_indexes]
        
        return FusionNode(nodes,node_value,node_type)
    
    def list_fusion_nodes(self,features:CommonFeatureSet,image_builder:AstImageBuilder):
        return list_fusion_node_state(features,image_builder)

    
    def list_fusion_dependency(self,features:CommonFeatureSet):
        fusion_state={}

        anchor_nodes=get_anchor_node_js(features)
        image_builder=AstImageBuilder(features.ast_object)


        for node in anchor_nodes:
            neighbors=self.get_anchor_neighbor(node,image_builder)
            for node in neighbors:
                if node not in fusion_state:
                    fusion_state[node]=set()
                    
                fusion_state[node].add(node)

        return fusion_state
    
    def get_anchor_nodes(self,features:CommonFeatureSet):
        nodes=[]
    
        for node in features.leave_nodes:
            if node.ast_type in self.anchor_type:
                nodes.append(node)
                
        return nodes
    
    def get_anchor_neighbor(self,node:AstNode,image_builder:AstImageBuilder):
        vision=self.anchor_visions.get(node.ast_type,0)
    
        neighbor_nodes=image_builder.get_node_neighbor_in(image_builder.node2index[node],vision)
        
        return neighbor_nodes